MODULE SRFSN_REGRID_MOD
CONTAINS
SUBROUTINE SRFSN_REGRID(KIDIA,KFDIA,KLON,KLEVSN,LLNOSNOW,&
                    PSSN,PRSN,PWSN,PTSN,PDSNNEW,&
                    PSSNNEW,PRSNNEW,PWSNNEW,PTSNNEW)

USE PARKIND1 , ONLY : JPIM, JPRB
USE YOMHOOK  , ONLY : LHOOK, DR_HOOK, JPHOOK
USE ABORT_SURF_MOD

! (C) Copyright 2015- ECMWF.
!
! This software is licensed under the terms of the Apache Licence Version 2.0
! which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
! In applying this licence, ECMWF does not waive the privileges and immunities
! granted to it by virtue of its status as an intergovernmental organisation
! nor does it submit to any jurisdiction.
!**** *SRFSN_REGRID* - Snow vertical re-grid
!     PURPOSE.
!     --------
!          THIS ROUTINE REGRIDS MULTI-LAYER SNOW FIELDS
!          ACCORDINGLY TO THE NEW VERTICAL THICKNESs COMPUTED
!          AT THE BEGINNING OF THE TIME STEP. 

!**   INTERFACE.
!     ----------
!          *SRFSN_REGRID* IS CALLED FROM *SRFSN_DRIVER*.

!     PARAMETER   DESCRIPTION                                    UNITS
!     ---------   -----------                                    -----

!     INPUT PARAMETERS (INTEGER):
!    *KIDIA*      START POINT
!    *KFDIA*      END POINT
!    *KLON*      NUMBER OF HORIZONTAL POINTS
!    *KLEVSN*      NUMBER OF VERTICAL SNOW LAYERS


!     INPUT PARAMETERS (REAL):
!    *PSSN*      SNOW MULTI-LAYER MASS
!    *PRSN*      SNOW MULTI-LAYER DENSITY
!    *PWSN*      SNOW MULTI-LAYER LIQUID WATER
!    *PTSN*      SNOW MULTI-LAYER TEMPERATURE
!    *PDSNNEW*   NEW MULTI-LAYER SNOW DEPTH TO BE REGRIDDED

!     INPUT PARAMETERS (LOGICAL):
!    *LLNOSNOW*   NOSNOW MASK (TRUE THEN POINT IS SNOW-FREE)

!     OUTPUT PARAMETERS (REAL):
!    *PSSNNEW*      SNOW MULTI-LAYER MASS REGRIDDED
!    *PRSNNEW*      SNOW MULTI-LAYER DENSITY REGRIDDED
!    *PWSNNEW*      SNOW MULTI-LAYER LIQUID WATER REGRIDDED
!    *PTSNNEW*      SNOW MULTI-LAYER TEMPERATURE REGRIDDED

!     METHOD.
!     -------
!          

!     EXTERNALS.
!     ----------
!          NONE.

!     REFERENCE.
!     ----------
!          

!     Modifications:
!     Original   E. Dutra      ECMWF     04/12/2015

!     ------------------------------------------------------------------

IMPLICIT NONE

! Declaration of arguments 
INTEGER(KIND=JPIM), INTENT(IN)   :: KIDIA
INTEGER(KIND=JPIM), INTENT(IN)   :: KFDIA
INTEGER(KIND=JPIM), INTENT(IN)   :: KLON
INTEGER(KIND=JPIM), INTENT(IN)   :: KLEVSN
LOGICAL           , INTENT(IN)   :: LLNOSNOW(:) 


REAL(KIND=JPRB), INTENT(IN) :: PSSN(:,:)
REAL(KIND=JPRB), INTENT(IN) :: PRSN(:,:) 
REAL(KIND=JPRB), INTENT(IN) :: PTSN(:,:) 
REAL(KIND=JPRB), INTENT(IN) :: PWSN(:,:)
REAL(KIND=JPRB), INTENT(IN) :: PDSNNEW(:,:) 

REAL(KIND=JPRB), INTENT(OUT) :: PSSNNEW(:,:)
REAL(KIND=JPRB), INTENT(OUT) :: PRSNNEW(:,:) 
REAL(KIND=JPRB), INTENT(OUT) :: PTSNNEW(:,:) 
REAL(KIND=JPRB), INTENT(OUT) :: PWSNNEW(:,:)

 
! Local variables
REAL(KIND=JPRB)  :: ZDSN(KLON,KLEVSN),ZHSN(KLON,KLEVSN),ZHSNNEW(KLON,KLEVSN),&
                    ZDSNG(KLON,0:KLEVSN),ZDSNNEWG(KLON,0:KLEVSN)
REAL(KIND=JPRB)  :: ZREMAP(KLEVSN,KLEVSN)

REAL(KIND=JPHOOK) :: ZHOOK_HANDLE
REAL(KIND=JPRB) :: ZEPSILON,ZIHCAP,&
                   ZTMPL1(KLEVSN),ZTMPL2(KLEVSN),ZTMP0,ZTMP1
INTEGER(KIND=JPRB) :: JL,JK,JKK,IBAD
REAL(KIND=JPRB) :: ZMULF
REAL(KIND=JPRB) :: ZRTT, ZRLMLT, ZRHOCI, ZRHOICE
LOGICAL         :: LLREMAP(KLON)
!    -----------------------------------------------------------------
IF (LHOOK) CALL DR_HOOK('SRFSN_REGRID_MOD:SRFSN_REGRID',0,ZHOOK_HANDLE)

!    -----------------------------------------------------------------

ZRTT=273.16_JPRB
ZRLMLT=2.8345E+6_JPRB-2.5008E+6_JPRB
ZRHOCI=2.05E6_JPRB
ZRHOICE=920._JPRB

ZIHCAP=ZRHOCI/ZRHOICE

ZMULF=10E5_JPRB
ZEPSILON=ZMULF*EPSILON(ZEPSILON)

!   0. pre-compute
  
  ! INPUT LAYER DEPTHS
ZDSN(KIDIA:KFDIA,1:KLEVSN) = PSSN(KIDIA:KFDIA,1:KLEVSN) / PRSN(KIDIA:KFDIA,1:KLEVSN)
  ! INPUT LAYER HEAT CONTENT 
ZHSN(KIDIA:KFDIA,1:KLEVSN) = ZIHCAP*(PSSN(KIDIA:KFDIA,1:KLEVSN))*&
                      (PTSN(KIDIA:KFDIA,1:KLEVSN)-ZRTT)-ZRLMLT*(PSSN(KIDIA:KFDIA,1:KLEVSN)-PWSN(KIDIA:KFDIA,1:KLEVSN))

 ! DEFINE VERTICAL GRID START AT 0 AT THE SNOW SURFACE IS POSITIVE DOWNWARDS
ZDSNG(KIDIA:KFDIA,0:KLEVSN) = 0._JPRB
ZDSNNEWG(KIDIA:KFDIA,0:KLEVSN) = 0._JPRB
DO JK=1,KLEVSN
  DO JL=KIDIA,KFDIA
    ZDSNG(JL,JK) = ZDSNG(JL,JK-1)+ZDSN(JL,JK)
    ZDSNNEWG(JL,JK) = ZDSNNEWG(JL,JK-1)+PDSNNEW(JL,JK)
  ENDDO
ENDDO

 ! SANITY CHECK:
DO JL=KIDIA,KFDIA
  IF ((ABS(ZDSNNEWG(JL,KLEVSN)-ZDSNG(JL,KLEVSN))>MAX(1._JPRB,ZDSNG(JL,KLEVSN))*ZEPSILON)) THEN
    WRITE(*,*) "TOTAL DEPTH OF INPUT ARRAY AND NEW DEPTH DO NOT MATCH"
    WRITE(*,*) ZDSNNEWG(JL,KLEVSN)-ZDSNG(JL,KLEVSN),ZDSNNEWG(JL,:),ZDSNG(JL,:)
    CALL ABORT_SURF("SRFSN_REGRID_MOD:SRFSN_REGRID: TOTAL DEPTH OF INPUT ARRAY AND NEW DEPTH DO NOT MATCH")
  ENDIF
ENDDO

LLREMAP(KIDIA:KFDIA) = .TRUE.
DO JL=KIDIA,KFDIA
  IF ( (LLNOSNOW(JL)) .OR. (SUM(PSSN(JL,:) - PWSN(JL,:)) == 0._JPRB) ) THEN
    LLREMAP(JL)=.FALSE.
  ENDIF
ENDDO

!   1. MAIN LOOP ON POINTS

DO JL=KIDIA,KFDIA

  IF ( LLREMAP(JL) ) THEN
    ! 1. FIND REMAPING COEF
    ZREMAP(:,:)=0._JPRB 
    
    DO JK=1,KLEVSN
      ZTMPL1(1:KLEVSN) = MAX(ZDSNNEWG(JL,JK-1),ZDSNG(JL,0:KLEVSN-1)) ! TOP BOUNDARY
      ZTMPL2(1:KLEVSN) = MIN(ZDSNNEWG(JL,JK),ZDSNG(JL,1:KLEVSN))     ! BOTTOM BOUNDARY 
      ZREMAP(JK,1:KLEVSN) = MAX(0._JPRB,ZTMPL2(1:KLEVSN)-ZTMPL1(1:KLEVSN))
  !     PRINT*,JK, ZREMAP(JK,:),SUM(ZREMAP(JK,:))
    ENDDO
    ! Sanity check if the remap matrix is correct, i.e. conservative 
    ! Here we relax the condiction to be 10*ZEPSILON  
    IF (ANY (ABS(SUM(ZREMAP(:,:),DIM=2)-PDSNNEW(JL,:))> MAX(1._JPRB,PDSNNEW(JL,:))*ZEPSILON*10._JPRB ) ) THEN
      write(*,*)JL,'ZREMAP',ABS(SUM(ZREMAP(:,:),DIM=2)-PDSNNEW(JL,:))> MAX(1._JPRB,PDSNNEW(JL,:))*ZEPSILON 
      DO JK=1,KLEVSN
         IF ( (ABS(SUM(ZREMAP(JK,:))-PDSNNEW(JL,JK))> MAX(1._JPRB,PDSNNEW(JL,JK))*ZEPSILON ) ) THEN
           write(*,*)JL,JK,'REMAP ERROR'
         ENDIF
         write(*,*)JL,JK,'REMAP ZREMAP',ZREMAP(JK,:)
         write(*,*)JL,JK,'REMAP 1',SUM(ZREMAP(JK,:)),PDSNNEW(JL,JK),MAX(1._JPRB,PDSNNEW(JL,JK))*ZEPSILON
         write(*,*)JL,JK,'REMAP 2',ABS(SUM(ZREMAP(JK,:))-PDSNNEW(JL,JK)),MAX(1._JPRB,PDSNNEW(JL,JK))*ZEPSILON
         write(*,*)JL,JK,'snowValOld',PSSN(JL,JK) , PRSN(JL,JK),ZDSN(JL,JK),PTSN(JL,JK),PWSN(JL,JK)
         write(*,*)JL,JK,'snowValNew',PDSNNEW(JL,JK),ZDSNNEWG(JL,JK)
         write(*,*)' ' 
         call flush(6)
      ENDDO
      CALL ABORT_SURF("SRFSN_REGRID_MOD:SRFSN_REGRID: REMAP MATRIX NOT CORRECT")
    ENDIF
    
    ! NORMALIZE BY LAYER DEPTHS
    DO JK=1,KLEVSN
      WHERE( ZREMAP(:,JK) > 0._JPRB)
        ZREMAP(:,JK) = ZREMAP(:,JK) / ZDSN(JL,JK)
      ENDWHERE
    ENDDO
    
    !* APPLY REMAPING
    PSSNNEW(JL,:) = MATMUL(ZREMAP,PSSN(JL,:))
    ZHSNNEW(JL,:) = MATMUL(ZREMAP,ZHSN(JL,:))
    
    ! REMAINING VALUES 
    DO JK=1,KLEVSN
      IF ( (PDSNNEW(JL,JK)>0._JPRB) ) THEN
        PRSNNEW(JL,JK) = PSSNNEW(JL,JK) / PDSNNEW(JL,JK)
        PTSNNEW(JL,JK) = MIN(ZRTT, ZRTT + (ZHSNNEW(JL,JK)+ZRLMLT*PSSNNEW(JL,JK))/&
                                 (ZIHCAP*PSSNNEW(JL,JK)))
        PWSNNEW(JL,JK) = MAX(0._JPRB, MIN(PSSNNEW(JL,JK), PSSNNEW(JL,JK) + ZHSNNEW(JL,JK)/ZRLMLT ))
      ELSE
        PWSNNEW(JL,JK) = 0._JPRB
        PTSNNEW(JL,JK) = PTSN(JL,JK)
        PRSNNEW(JL,JK) = PRSN(JL,JK)
      ENDIF
    ENDDO

  ELSE
    PSSNNEW(JL,:) = PSSN(JL,:)
    PTSNNEW(JL,:) = PTSN(JL,:)
    PRSNNEW(JL,:) = PRSN(JL,:)
    PWSNNEW(JL,:) = PWSN(JL,:)
    ZHSNNEW(JL,:) = ZHSN(JL,:)
  ENDIF 
  
ENDDO
! PRINT*,'RE,ZDSN',ZDSN
! PRINT*,'RE,PDSNNEW',PDSNNEW
! PRINT*,'RE,ZDSNG',ZDSNG
! PRINT*,'RE,ZDSNNEWG',ZDSNNEWG
! PRINT*,'RE,PRSN',PRSN
! PRINT*,'RE,PRSNNEW',PRSNNEW
! PRINT*,'RE,PSSN',PSSN
! PRINT*,'RE,PSSNNEW',PSSNNEW
! PRINT*,'RE,PWSN',PWSN
! PRINT*,'RE,PWSNNEW',PWSNNEW
! PRINT*,'RE,PTSN',PTSN
! PRINT*,'RE,PTSNNEW',PTSNNEW

!! CHECKS
! INPUT LAYER HEAT CONTENT 
ZHSNNEW(KIDIA:KFDIA,:) = ZIHCAP*(PSSNNEW(KIDIA:KFDIA,:))*&
                      (PTSNNEW(KIDIA:KFDIA,:)-ZRTT)-ZRLMLT*(PSSNNEW(KIDIA:KFDIA,:)-PWSNNEW(KIDIA:KFDIA,:))


DO JL=KIDIA,KFDIA
    IBAD=0
    ZTMP1 = SUM(PSSN(JL,:))
    IF (ZTMP1 > ZEPSILON ) THEN            
      ZTMP0 = (SUM(PSSNNEW(JL,:))-SUM(PSSN(JL,:)))/ZTMP1
    !   PRINT*,'WATER',ZTMP0
      IF ( ABS(ZTMP0) > ZEPSILON ) THEN
        WRITE(*,*) JL,ZTMP0,SUM(PSSNNEW(JL,:)),ZTMP1
        WRITE(*,*) " MASS NOT CONSERVED IN SNOW REGRID!"
        IBAD=1
      ENDIF
      ZTMP0 = (SUM(ZHSN(JL,:))-SUM(ZHSNNEW(JL,:)))/MIN(-ZEPSILON, SUM(ZHSNNEW(JL,:)))
    !   PRINT*,'EN1',ZTMP0
      IF ( ABS(ZTMP0) > ZEPSILON ) THEN
        WRITE(*,*) JL,ZTMP0,SUM(ZHSN(JL,:)),SUM(ZHSNNEW(JL,:))
        WRITE(*,*) " ENERGY NOT CONSERVED IN SNOW REGRID!"
        IBAD=1
      ENDIF
      IF ( IBAD>0) THEN
        WRITE(*,*)'PDSN',ZDSN(JL,:)
        WRITE(*,*)'PDSNNEW',PDSNNEW(JL,:)
        WRITE(*,*)'PSSN',PSSN(JL,:)
        WRITE(*,*)'PSSNNEW',PSSNNEW(JL,:)
        WRITE(*,*)'PWSN',PWSN(JL,:)
        WRITE(*,*)'PWSNNEW',PWSNNEW(JL,:)
        WRITE(*,*)'PRSN',PRSN(JL,:)
        WRITE(*,*)'PRSNNEW',PRSNNEW(JL,:)
        WRITE(*,*)'PTSN',PTSN(JL,:)
        WRITE(*,*)'PTSNNEW',PTSNNEW(JL,:)
        WRITE(*,*)'ZHSN',ZHSN(JL,:)
        WRITE(*,*)'ZHSNNEW',ZHSNNEW(JL,:)  
        IF ( IBAD == 1) THEN
          CALL ABORT_SURF('SNOW REGRID ERROR!')
        ENDIF
      ENDIF
    ENDIF
ENDDO

!    -----------------------------------------------------------------
IF (LHOOK) CALL DR_HOOK('SRFSN_REGRID_MOD:SRFSN_REGRID',1,ZHOOK_HANDLE)

END SUBROUTINE SRFSN_REGRID
END MODULE SRFSN_REGRID_MOD
